package com.dkd.manage.task.allocator.algorithm;

import com.dkd.manage.domain.Emp;
import lombok.extern.slf4j.Slf4j;

import java.util.*;
import java.util.function.Function;

@Slf4j
public class AntColonyOptimization {
    private int antCount;
    private int iterations;
    private double alpha; // 信息素重要性
    private double beta;  // 启发式信息重要性
    private double evaporation;

    public AntColonyOptimization(int antCount, int iterations, double alpha, double beta, double evaporation) {
        this.antCount = antCount;
        this.iterations = iterations;
        this.alpha = alpha;
        this.beta = beta;
        this.evaporation = evaporation;
    }

    public Map<Long, Double> optimize(List<Emp> candidates, Function<Emp, Double> heuristicFunction) {
        log.info("开始蚁群优化，候选人数: {}", candidates.size());

        Map<Long, Double> results = new HashMap<>();
        Map<Long, Double> pheromone = initializePheromone(candidates);

        for (int iter = 0; iter < iterations; iter++) {
            Map<Long, Double> deltaPheromone = new HashMap<>();

            // 每只蚂蚁构建解
            for (int ant = 0; ant < antCount; ant++) {
                Emp selected = constructSolution(candidates, pheromone, heuristicFunction);

                // 更新信息素
                double quality = heuristicFunction.apply(selected);
                deltaPheromone.merge(selected.getId(), quality, Double::sum);
            }

            // 更新全局信息素
            updatePheromone(pheromone, deltaPheromone);

            log.debug("蚁群优化第 {} 代完成", iter + 1);
        }

        // 生成最终结果
        for (Emp emp : candidates) {
            double baseScore = heuristicFunction.apply(emp);
            double pheromoneLevel = pheromone.getOrDefault(emp.getId(), 1.0);
            double optimizedScore = baseScore * (0.6 + 0.8 * pheromoneLevel);
            results.put(emp.getId(), optimizedScore);
        }

        log.info("蚁群优化完成，生成 {} 个结果", results.size());
        return results;
    }

    private Map<Long, Double> initializePheromone(List<Emp> candidates) {
        Map<Long, Double> pheromone = new HashMap<>();
        for (Emp emp : candidates) {
            pheromone.put(emp.getId(), 1.0); // 初始信息素
        }
        return pheromone;
    }

    private Emp constructSolution(List<Emp> candidates, Map<Long, Double> pheromone,
                                  Function<Emp, Double> heuristicFunction) {
        Random random = new Random();
        List<Double> probabilities = new ArrayList<>();

        // 计算选择概率
        for (Emp emp : candidates) {
            double pheromoneLevel = Math.pow(pheromone.get(emp.getId()), alpha);
            double heuristic = Math.pow(heuristicFunction.apply(emp), beta);
            probabilities.add(pheromoneLevel * heuristic);
        }

        // 轮盘赌选择
        double total = probabilities.stream().mapToDouble(Double::doubleValue).sum();
        double randomValue = random.nextDouble() * total;

        double cumulative = 0.0;
        for (int i = 0; i < candidates.size(); i++) {
            cumulative += probabilities.get(i);
            if (cumulative >= randomValue) {
                return candidates.get(i);
            }
        }

        return candidates.get(candidates.size() - 1);
    }

    private void updatePheromone(Map<Long, Double> pheromone, Map<Long, Double> deltaPheromone) {
        // 信息素蒸发
        for (Long empId : pheromone.keySet()) {
            pheromone.put(empId, pheromone.get(empId) * (1 - evaporation));
        }

        // 信息素增加
        for (Map.Entry<Long, Double> entry : deltaPheromone.entrySet()) {
            pheromone.merge(entry.getKey(), entry.getValue(), Double::sum);
        }
    }
}